package netpool

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"time"
)

type idleConn struct {
	c              any
	lastActiveTime time.Time
}

type connReq struct {
	connChan chan any
}

type Pool struct {
	// 空闲队列
	idleConns   chan *idleConn
	maxIdleTime time.Duration
	maxIdleCnt  int32

	factory func() (any, error)
	mu      sync.Mutex
	cnt     int32
	maxCnt  int32

	// 等待队列
	reqQueue []connReq

	close func(any) error
}

func NewPool(initCnt int32, maxCnt int32, maxIdleCnt int32, maxIdleTime time.Duration, factory func() (any, error), close func(any) error) (Pooler, error) {
	if initCnt > maxCnt {
		return nil, errors.New("pool: 初始连接数量不能大于最大连接数量")
	}
	if factory == nil || close == nil {
		return nil, errors.New("pool: factory or close is nil")
	}

	idlesConns := make(chan *idleConn, maxIdleCnt)
	for i := int32(0); i < initCnt; i++ {
		conn, err := factory()
		if err != nil {
			return nil, err
		}
		idlesConns <- &idleConn{c: conn, lastActiveTime: time.Now()}
	}
	res := &Pool{
		idleConns:   idlesConns,
		maxIdleCnt:  maxIdleCnt,
		maxCnt:      maxCnt,
		maxIdleTime: maxIdleTime,
		factory:     factory,
		close:       close,
		cnt:         initCnt,
	}
	return res, nil
}

func (p *Pool) Get(ctx context.Context) (any, error) {
	if p.idleConns == nil {
		return nil, errors.New("pool: pool is closed")
	}

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case ic := <-p.idleConns:
			if ic.lastActiveTime.Add(p.maxIdleTime).Before(time.Now()) {
				p.Close(ic.c)
				continue
			}
			return ic.c, nil
		default:
			p.mu.Lock()
			if p.cnt >= p.maxCnt {
				req := connReq{connChan: make(chan any, 1)}
				p.reqQueue = append(p.reqQueue, req)
				p.mu.Unlock()
				select {
				case <-ctx.Done():
					// 虽然超时，但等待队列中的等待请求没有移除，后续还可能被唤醒，
					// 开一个 goroutine 来处理，把收到的连接再调用 Put 方法放到池中
					go func() {
						conn := <-req.connChan
						_ = p.Put(context.Background(), conn)
					}()
					return nil, ctx.Err()
				case conn := <-req.connChan:
					return conn, nil
				}
			}

			// 获取连接的过程耗时大，最好先解锁，获取连接失败再用原子操作把连接数量-1
			p.cnt++
			p.mu.Unlock()

			conn, err := p.factory()
			if err != nil {
				atomic.AddInt32(&p.cnt, -1)
				return nil, err
			}

			return conn, nil
		}
	}
}

func (p *Pool) Put(ctx context.Context, conn any) error {
	if p.idleConns == nil {
		return errors.New("pool: pool is closed")
	}

	p.mu.Lock()
	if len(p.reqQueue) > 0 {
		req := p.reqQueue[0]
		p.reqQueue = p.reqQueue[1:]
		p.mu.Unlock()
		req.connChan <- conn
		return nil
	}
	p.mu.Unlock()

	ic := &idleConn{
		c:              conn,
		lastActiveTime: time.Now(),
	}
	select {
	case p.idleConns <- ic:
	default:
		p.Close(conn)
	}
	return nil
}

func (p *Pool) Close(c any) error {
	if c == nil || p.close == nil {
		return nil
	}

	if err := p.close(c); err == nil {
		atomic.AddInt32(&p.cnt, -1)
	}

	return nil
}

func (p *Pool) Release() {
	p.mu.Lock()
	idleConns := p.idleConns
	p.idleConns = nil
	p.factory = nil
	closeFun := p.close
	p.close = nil
	p.cnt = 0
	p.mu.Unlock()

	if idleConns == nil {
		return
	}

	close(idleConns)
	for idleConn := range idleConns {
		closeFun(idleConn.c)
	}
}

func (p *Pool) Len() int {
	return len(p.idleConns)
}
